"""
子问题求解器

提供各种优化子问题的求解方法
"""

import numpy as np
import cvxpy as cp
from typing import Callable, Optional, Dict, Any, Tuple


def solve_local_subproblem(objective: Callable,
                          gradient: Optional[Callable],
                          x0: np.ndarray,
                          constraints: Dict[str, Any],
                          method: str = 'cvxpy') -> np.ndarray:
    """
    求解局部优化子问题
    
    参数:
        objective: 目标函数
        gradient: 梯度函数（可选）
        x0: 初始点
        constraints: 约束描述
        method: 求解方法 ('cvxpy', 'gradient', 'proximal')
        
    返回:
        最优解
    """
    if method == 'cvxpy':
        return _solve_with_cvxpy(objective, x0, constraints)
    elif method == 'gradient':
        return _solve_with_gradient(objective, gradient, x0, constraints)
    elif method == 'proximal':
        return _solve_with_proximal(objective, gradient, x0, constraints)
    else:
        raise ValueError(f"未知的求解方法: {method}")
        

def _solve_with_cvxpy(objective: Callable,
                     x0: np.ndarray,
                     constraints: Dict[str, Any]) -> np.ndarray:
    """使用CVXPY求解"""
    n = len(x0)
    x = cp.Variable(n)
    
    # 目标函数
    obj = objective(x)
    
    # 构建约束
    cvx_constraints = []
    
    if 'type' in constraints:
        set_type = constraints['type']
        
        if set_type == 'box':
            lower = constraints['lower']
            upper = constraints['upper']
            cvx_constraints.append(x >= lower)
            cvx_constraints.append(x <= upper)
            
        elif set_type == 'ball':
            center = constraints['center']
            radius = constraints['radius']
            cvx_constraints.append(cp.norm(x - center) <= radius)
            
        elif set_type == 'simplex':
            cvx_constraints.append(x >= 0)
            cvx_constraints.append(cp.sum(x) <= 1)
            
    # 求解
    prob = cp.Problem(cp.Minimize(obj), cvx_constraints)
    try:
        prob.solve(solver=cp.OSQP, verbose=False)
        if prob.status == cp.OPTIMAL:
            return x.value
        else:
            return x0
    except:
        return x0
        

def _solve_with_gradient(objective: Callable,
                        gradient: Callable,
                        x0: np.ndarray,
                        constraints: Dict[str, Any],
                        max_iter: int = 100,
                        tol: float = 1e-6) -> np.ndarray:
    """使用投影梯度法求解"""
    x = x0.copy()
    
    for k in range(max_iter):
        # 计算梯度
        grad = gradient(x)
        
        # 自适应步长
        step_size = 0.1 / (1 + k)
        
        # 梯度步
        x_new = x - step_size * grad
        
        # 投影到约束集
        x_new = project_to_set(x_new, constraints)
        
        # 检查收敛
        if np.linalg.norm(x_new - x) < tol:
            break
            
        x = x_new
        
    return x
    

def _solve_with_proximal(objective: Callable,
                        gradient: Optional[Callable],
                        x0: np.ndarray,
                        constraints: Dict[str, Any],
                        rho: float = 1.0,
                        max_iter: int = 100) -> np.ndarray:
    """使用近端梯度法求解"""
    x = x0.copy()
    
    for k in range(max_iter):
        if gradient is not None:
            # 近端梯度步
            grad = gradient(x)
            step_size = 1.0 / (rho * (1 + k))
            x_grad = x - step_size * grad
        else:
            x_grad = x
            
        # 近端算子（投影）
        x_new = project_to_set(x_grad, constraints)
        
        # 检查收敛
        if np.linalg.norm(x_new - x) < 1e-6:
            break
            
        x = x_new
        
    return x
    

def project_to_set(x: np.ndarray, constraints: Dict[str, Any]) -> np.ndarray:
    """投影到约束集"""
    if 'type' not in constraints:
        return x
        
    set_type = constraints['type']
    
    if set_type == 'box':
        lower = np.array(constraints['lower'])
        upper = np.array(constraints['upper'])
        return np.clip(x, lower, upper)
        
    elif set_type == 'ball':
        center = np.array(constraints['center'])
        radius = constraints['radius']
        diff = x - center
        norm = np.linalg.norm(diff)
        if norm <= radius:
            return x
        else:
            return center + radius * diff / norm
            
    elif set_type == 'simplex':
        return project_to_simplex(x)
        
    else:
        return x
        

def project_to_simplex(x: np.ndarray) -> np.ndarray:
    """投影到单纯形 {x: x >= 0, sum(x) <= 1}"""
    n = len(x)
    if np.all(x >= 0) and np.sum(x) <= 1:
        return x
        
    # 排序算法
    u = np.sort(x)[::-1]
    cssv = np.cumsum(u)
    rho = np.where(u > (cssv - 1) / np.arange(1, n + 1))[0][-1]
    theta = (cssv[rho] - 1) / (rho + 1)
    return np.maximum(x - theta, 0)
    

def solve_quadratic_program(Q: np.ndarray, c: np.ndarray,
                           A_ineq: Optional[np.ndarray] = None,
                           b_ineq: Optional[np.ndarray] = None,
                           A_eq: Optional[np.ndarray] = None,
                           b_eq: Optional[np.ndarray] = None,
                           lb: Optional[np.ndarray] = None,
                           ub: Optional[np.ndarray] = None) -> Tuple[np.ndarray, float]:
    """
    求解二次规划问题
    
    minimize    (1/2) x^T Q x + c^T x
    subject to  A_ineq x <= b_ineq
               A_eq x = b_eq
               lb <= x <= ub
    """
    n = len(c)
    x = cp.Variable(n)
    
    # 目标函数
    objective = 0.5 * cp.quad_form(x, Q) + c.T @ x
    
    # 约束
    constraints = []
    
    if A_ineq is not None and b_ineq is not None:
        constraints.append(A_ineq @ x <= b_ineq)
        
    if A_eq is not None and b_eq is not None:
        constraints.append(A_eq @ x == b_eq)
        
    if lb is not None:
        constraints.append(x >= lb)
        
    if ub is not None:
        constraints.append(x <= ub)
        
    # 求解
    prob = cp.Problem(cp.Minimize(objective), constraints)
    prob.solve(solver=cp.OSQP, verbose=False)
    
    if prob.status == cp.OPTIMAL:
        return x.value, prob.value
    else:
        raise ValueError(f"二次规划求解失败: {prob.status}")
        

def proximal_operator(x: np.ndarray, f: Callable, t: float,
                     method: str = 'gradient') -> np.ndarray:
    """
    计算近端算子
    
    prox_{tf}(x) = argmin_y { f(y) + (1/2t)||y - x||^2 }
    """
    if method == 'gradient':
        # 使用梯度下降
        y = x.copy()
        for _ in range(100):
            if hasattr(f, 'gradient'):
                grad = f.gradient(y)
            else:
                # 数值梯度
                eps = 1e-6
                grad = np.zeros_like(y)
                for i in range(len(y)):
                    y_plus = y.copy()
                    y_minus = y.copy()
                    y_plus[i] += eps
                    y_minus[i] -= eps
                    grad[i] = (f(y_plus) - f(y_minus)) / (2 * eps)
                    
            # 近端梯度步
            y_new = y - t * (grad + (y - x) / t)
            
            if np.linalg.norm(y_new - y) < 1e-8:
                break
                
            y = y_new
            
        return y
        
    else:
        raise ValueError(f"未知的近端算子计算方法: {method}") 